import numpy as np
from qutip import Qobj, sigmax, sigmay, sigmaz, mesolve, Options
import scipy.optimize

# Define physical constants
h = 6.626e-34  # Planck's constant (J·s)
c = 3.00e8     # Speed of light (m/s)
hbar = h / (2 * np.pi)

# Parameter ranges
omegas = np.array([0.05e6, 0.1e6, 0.2e6, 0.5e6])  # Rabi frequencies (MHz)
gammas = np.array([0.05e6, 0.1e6, 0.2e6])          # Decoherence rates (MHz)
wavelengths = np.array([630e-9])                   # Single wavelength for simplicity (m)
omega_0 = 2 * np.pi * c / wavelengths[0]           # Transition frequency (Hz)

# TLS Hamiltonian
sigma_x, sigma_y, sigma_z = sigmax(), sigmay(), sigmaz()
H0 = 0.5 * omega_0 * sigma_z
H1 = 0.5 * sigmax()

# Time parameters
tlist = np.linspace(0, 50e-6, 500)  # Time array (µs)

# Function to solve dynamics and extract steady-state
def get_steady_state(omega, gamma):
    H = H0 + 0.5 * omega * H1
    psi0 = Qobj([[1], [0]])
    c_ops = [np.sqrt(g) * s for g, s in zip([gamma], [sigma_x, sigma_y, sigma_z])]
    result = mesolve(H, psi0, tlist, c_ops=c_ops, e_ops=[sigmaz()], options=Options(store_states=True))
    steady_pop = 0.5 * (1 - result.expect[0][-1])  # Last time point
    return steady_pop

# Function to fit coherence decay (exponential decay model)
def exp_decay(t, A, tau):
    return A * np.exp(-t / tau)

def fit_coherence_decay(omega, gamma):
    H = H0 + 0.5 * omega * H1
    psi0 = Qobj([[1], [0]])
    c_ops = [np.sqrt(g) * s for g, s in zip([gamma], [sigma_x, sigma_y, sigma_z])]
    result = mesolve(H, psi0, tlist, c_ops=c_ops, e_ops=[sigmax(), sigmay()], options=Options(store_states=True))
    coherence = np.sqrt(result.expect[0]**2 + result.expect[1]**2)
    popt, _ = scipy.optimize.curve_fit(exp_decay, tlist, coherence, p0=[1.0, 1e-6])
    return popt[0], popt[1]  # Amplitude A, decay time tau (s)

# Generate derived datasets
steady_state_data = []
coherence_fit_data = []
sample_excerpt_data = []

for omega in omegas:
    for gamma in gammas:
        # Steady-state population
        steady_pop = get_steady_state(omega, gamma)
        steady_state_data.append([omega * 1e-6, gamma * 1e-6, steady_pop])
        
        # Coherence decay fit parameters
        A, tau = fit_coherence_decay(omega, gamma)
        coherence_fit_data.append([omega * 1e-6, gamma * 1e-6, A, tau * 1e6])  # tau in µs
        
        # Sample excerpt (first 10 time points)
        H = H0 + 0.5 * omega * H1
        psi0 = Qobj([[1], [0]])
        c_ops = [np.sqrt(g) * s for g, s in zip([gamma], [sigma_x, sigma_y, sigma_z])]
        result = mesolve(H, psi0, tlist[:10], c_ops=c_ops, e_ops=[sigmaz()])
        excited_pop = 0.5 * (1 - result.expect[0])
        for t, pop in zip(tlist[:10] * 1e6, excited_pop):  # Time in µs
            sample_excerpt_data.append([t, pop, omega * 1e-6, gamma * 1e-6])

# Save to CSV (example for steady-state, adapt for others)
import pandas as pd
steady_df = pd.DataFrame(steady_state_data, columns=['Omega (MHz)', 'Gamma (MHz)', 'Steady_State_Excited_Population'])
steady_df.to_csv('tls_steady_state_populations.csv', index=False)